import abc
from typing import Any, Optional
from langchain.tools import BaseTool
import logging
import sys
import pandas as pd
from utils.configFinRAG import sql_examples_path  # 假设配置文件中有sql_examples_path
from FinSQL_01_generate import generate_sql
from FinSQL_02_query import query_db
from FinSQL_03_answer_from_SQL import generate_answer
from utils.instances import TOKENIZER, LLM

logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))


class FinSQLRAG(BaseTool, abc.ABC):
    name = "查询金融数据库"
    description = "当被问到金融查询相关的问题时，会去金融数据库检索结果"

    def __init__(self, db_path: str):
        super().__init__()
        self.db_path = db_path
        self.example_lists = {
            'questions': [],
            'sqls': [],
            'infos': [],
            'fa': [],
            'tokens': []
        }
        self.load_example_data()

    def load_example_data(self):
        """加载示例数据"""
        examples = pd.read_csv(sql_examples_path, delimiter=",", header=0)
        for _, row in examples.iterrows():
            self.example_lists['questions'].append(row['问题'])
            self.example_lists['sqls'].append(row['SQL'])
            self.example_lists['infos'].append(row['资料'])
            self.example_lists['fa'].append(row['FA'])
            tokens = TOKENIZER(row['问题'])['input_ids']
            self.example_lists['tokens'].append(tokens)

    async def _arun(self, *args: Any, **kwargs: Any) -> Any:
        # 用例中没有用到 arun 不予具体实现
        pass

    def _run(self, para: str) -> str:
        query = para
        try:
            result_prompt, sql = generate_sql(query, LLM, self.example_lists['questions'], self.example_lists['sqls'],
                                              self.example_lists['tokens'])
            logging.info(f"Generated SQL query is: {sql}")

            with sqlite3.connect(self.db_path) as conn:
                cs = conn.cursor()
                success_flag, exc_result = query_db(sql, cs)

            if not success_flag:
                raise Exception("Database query failed.")

            answer = generate_answer(query, exc_result, LLM, self.example_lists['questions'],
                                     self.example_lists['infos'],
                                     self.example_lists['fa'], self.example_lists['tokens'])
            return answer
        except Exception as e:
            logging.error(f"Retriever Error: {e}")
            return 'FinSQLRAG处理异常！'


if __name__ == "__main__":
    db_path = '/path/to/your/database.db'  # 根据实际情况修改路径
    tool = FinSQLRAG(db_path=db_path)
    result = tool._run(
        "请帮我计算，在20210715，中信行业分类划分的一级行业为消费者服务行业中，涨跌幅最大股票的股票代码是？涨跌幅是多少？百分数保留两位小数。股票涨跌幅定义为：（收盘价 - 前一日收盘价 / 前一日收盘价）* 100%。")
    print(result)